set type double

// Recover beliefs from reports given a utility and probability weighting function
program define beliefs_recovery


	// Utility and Probability Weighting functions need to be defined
	// The strlower function is used to sanatize the input
	local ufunc = strlower("$ufunc")
	local pfunc = strlower("$pfunc")

	// Reports begin with 'r_'
	global rstub "r_"
	// Make sure there are no missing reports
	forval i = 1 / $nbin {
		capture generate $rstub`i' = 0
	}
	order $rstub*, sequential

	// Beliefs begin with 'b_'
	global bstub "b_"

	// Depending on the kind of utility and probability weighting functions used,
	// there will be a different number of arguments to this program
	local num_args = 0

	// Currently supporting:
	// CRRA       (crra)  - x^(1 - r) / (1 - r)
	// Power      (power) - x^r
	// CARA       (cara)  - exp(-k * x)
	// Expo-Power (expo)  - [1 - exp(-a * x^(1 - s))] / a
	if "`ufunc'" == "crra" | "`ufunc'" == "power" | "`ufunc'" == "cara" {
		local num_args = `num_args' + 1
	}
	else if "`ufunc'" == "expo" {
		local num_args = `num_args' + 2
	}
	else	{
		display "You are using an unsupported utility function for beliefs recovery"
		error 6
	}

	// This is what's called a 'linked list'. Stata doesn't do them natively,
	// they need to be constructed with local macros
	forvalues i = 1/`num_args' {
		local uargs "`uargs' par`i'"
	}
	// The number of utility arguments
	local nuargs = `num_args'

	// Currently supporting 2-parameter Prelec, Power, and  Inverse-S.
	// The inverse-s function is provided via numeric approximation as it's
	// inverse does not have an analytical solution.
	if "`pfunc'" == "eut" | "`pfunc'" == "" {
		// Nothing to do here
	}
	else if "`pfunc'" == "power" {
		local num_args = `num_args' + 1
	}
	else if "`pfunc'" == "inverse_s" {
		local num_args = `num_args' + 1
	}
	else if "`pfunc'" == "prelec2" {
		local num_args = `num_args' + 2
	}
	else if "`pfunc'" != "" {
		display "You are using an unsupported probability weighting function for beliefs recovery"
		error 6
	}
	local start = `nuargs' + 1
	forvalues i = `start'/`num_args' {
		local pargs "`pargs' par`i'"
	}

	// uargs - The utility function parameters
	// pargs - The probability weighting function parameters
	args `uargs' `pargs'

	// Get the *actual* variables referenced by uargs and pargs
	// This is resolving the linked list
	foreach val of local uargs {
		local upars "`upars' ``val''"
	}
	foreach val of local pargs {
		local ppars "`ppars' ``val''"
	}

	// Now actually start recovering beliefs
	// Make sure reports sum to 1
	tempvar report_total
	qui egen `report_total' = rowtotal($rstub*)
	foreach report of var $rstub* {
		qui replace `report' = `report' / `report_total'
	}

	// Part of the qsr is the sum of squarred reports
	// We can efficiently get this done now
	tempvar report_sq
	qui gen double `report_sq' = 0
	foreach report of var $rstub* {
		qui replace `report_sq' = `report_sq' + `report'^2
	}

	// Translate the reports to QSR outcomes
	forval i = 1 / $nbin {
		tempvar outcome_`i'
		qui gen double `outcome_`i'' = $alpha + 2 * $beta * $rstub`i' - $beta * `report_sq'
	}

	// Trying to be mindful of RAM usage. Some applications have a _lot_ of
	// observations, and so we'll drop variables after they are no longer useful
	drop `report_sq'

	// Get the first derivative of the utility function
	forval i = 1 / $nbin {
		tempvar out_prime_`i'
		qui gen double `out_prime_`i'' = 0
		`ufunc'_prime `out_prime_`i'' `outcome_`i'' `upars'
	}

	// Get the decision weights implied by these outcomes
	forval i = 1 / $nbin {
		qui gen double dweight_`i' = $rstub`i' / `out_prime_`i''
		qui replace dweight_`i' = 0 if $rstub`i' == 0
		// No longer need this outcome prime number
		drop `out_prime_`i''
	}

	tempvar dw_total
	qui egen `dw_total' = rowtotal(dweight_*),

	forval i = 1 / $nbin {
		// Ensure the decision weights sum to one
		qui replace dweight_`i' = dweight_`i' / `dw_total'
		// Generate the variables that will hold the recovered beliefs
		capture generate $bstub`i' = 0
	}

	// No longer need the rowtotal of decision weights
	drop `dw_total'

	if "`pfunc'" == "eut" | "`pfunc'" == "" {
		// If we're in EUT we can skip the whole decumulating decision weights
		// thing
		forval i = 1 / $nbin {
			qui replace $bstub`i' = dweight_`i'
			qui drop dweight_`i'
		}
	}
	else {
		// RDU requires some serious extra effort provided below
		decumulate_dw `ppars'
		// We no longer need the dweight_* variables
		forval i = 1 / $nbin {
			qui drop dweight_`i'
		}
	}

	// Ensure that beliefs sum to 1
	tempvar bel_checker
	qui egen `bel_checker' = rowtotal($bstub*)
	forval i = 1 / $nbin {
		qui replace $bstub`i' = $bstub`i' / `bel_checker'
		qui replace $bstub`i' = 0 if $bstub`i' == .
	}

end

// Decumulate decision weights if probability weighting is not linear
program define decumulate_dw
	// Get the probability weighting parameters
	local counter = 1
	while 1 {
		if "``counter''" != "" {
			local ppars "`ppars' ``counter''"
		}
		else {
			continue, break
		}
		local ++counter
	}

	// Get the probability weighting funciton
	local pfunc = strlower("$pfunc")

	// Loop through all the decision weights and save their order
	forval i = 1 / $nbin {
		tempvar order`i'
		qui gen `order`i'' = 1
		forval j = 1 / $nbin {
			qui replace `order`i'' = `order`i'' + 1 if dweight_`j' > dweight_`i'
		}
	}

	// The above sets the highest ranked outcome as 1 and the lowest as equal
	// to the number of bins. If a bin has the same report as one or more other
	// bins, then they will share the rank.

	// The below increments the order if it is shared with one or more bins.
	forval i = 1 / $nbin {
		local ii = `i' + 1
		forval j = `ii' / $nbin {
			qui replace `order`i'' = `order`i'' + 1 if dweight_`j' == dweight_`i'
		}
	}

	// Sort the decision weights based on their order
	forval i = 1 / $nbin {
		tempvar dwr_`i'
		qui gen double `dwr_`i'' = 0
		forval j = 1 / $nbin {
			qui replace `dwr_`i'' = dweight_`j' if `order`j'' == `i'
		}

	}

	// When reports are equal for one or more bins, they must have the same
	// probability of occuring. But this presents a problem with probability
	// weighting which requires a strict ordering of outcomes. Thus, we
	// allocate the decision weights
	forval i = 1 / $nbin {
		tempvar dw_`i'
		local vals "`vals' `dw_`i''"
		qui gen double `dw_`i'' = `dwr_`i''
	}

	forval i = 1 / $nbin {
		local ii = `i' + 1
		forval j = `ii' / $nbin {
			qui replace `dw_`i'' =  `dw_`i'' + `dw_`j'' if `dwr_`i'' == `dwr_`j''
			qui replace `dw_`j'' =  0                   if `dwr_`i'' == `dwr_`j''
		}
	}

	// Cummulatively sum the decision weights
	forval i = 1 / $nbin {
		// Cumulative dw variable
		tempvar cdw_`i'
		qui gen double `cdw_`i'' = `dw_`i''
		// If we're on the highest ranked dw, nothing to cummulate
		if `i' == 1 {
			continue
		}
		// Add the current dw to the previous cummulative dw
		local j = `i' - 1
		qui replace `cdw_`i'' = `cdw_`i'' + `cdw_`j''
		// Make sure we don't have any bins with cummulative probs > 1
		qui replace `cdw_`i'' = 1 if `cdw_`i'' > 1
		// Store the final bin so we can force it to equal 1
		local final_bin "`cdw_`i''"
	}

	// Because of computer limitations, this final bin somtimes does not
	// perfectly sum to 1. It's within the noise level of a float precision
	// variable, but not within a double. Force it to be equal to 1.
	// See Gould (2006) "Mata Matters: Precision" for a discussion about what's
	// happening and why it matters
	qui replace `final_bin' = 1

	// Apply the inverse probability weighting function to the cummulated decision weights
	forval i = 1 / $nbin {
		`pfunc'_inverse `cdw_`i'' `ppars'
	}

	// Decumulate the now cummulative probabilities
	forval i = 1 / $nbin {
		tempvar bel_`i'
		qui gen double `bel_`i'' = 0
		// If we're on the highest ranked dw, nothing to decumulate
		if `i' == 1 {
			qui replace `bel_`i'' = `cdw_`i''
			continue
		}
		// Add the current dw to the previous cummulative dw
		local j = `i' - 1
		qui replace `bel_`i'' = `cdw_`i'' - `cdw_`j''
	}

	// Unsort the decumulated probabilities
	forval i = 1 / $nbin {
		forval j = 1 / $nbin {
			// bstub`i' is generated one program up. it isn't a tempvar
			qui replace $bstub`i' = `bel_`j'' if `order`i'' == `j'
		}
	}

	// For bins that had the same reports as other bins, split the reports
	forval i = 1 / $nbin {
		tempvar split_bel_`i'
		qui gen `split_bel_`i'' = 0
		forval j = 1 / $nbin {
			qui replace `split_bel_`i'' = `split_bel_`i'' + 1 if $rstub`i' == $rstub`j'
		}
	}

	// Divide the grouped beliefs
	forval i = 1 / $nbin {
		local ci "$rstub`i'"
		qui replace $bstub`i' = $bstub`i' / `split_bel_`i''
	}

	// Loop through the variables in reverse order because when grouped,
	// the last value got the summed probabilities
	forval i = $nbin(-1)1 {
		forval j = $nbin(-1)1 {
			qui replace $bstub`j' = $bstub`i' if $rstub`i' == $rstub`j'
		}
	}
end

// Inverse of Prelec (1998) probability weighting function
program define prelec2_inverse
	args var eta phi

	// Make a tempvar so we can replace 0s and 1s later
	tempvar dw
	qui gen `dw' = `var'

	// Do the inverse
	qui replace `dw' = log(`dw') / -`phi'
	qui replace `dw' = `dw' ^ (1 / `eta')
	qui replace `dw' = exp(-`dw')
	qui replace `dw' = 1 if `dw' == .

	// Hard replace 0s with 0s and 1s with 1s
	qui replace `dw' = 0 if `var' == 0
	qui replace `dw' = 1 if `var' == 1

	// Overwrite the initial value
	qui replace `var' = `dw'
end

// Inverse of Power, Quiggan (1998), probability weighting function
program define power_inverse
	args var gamma

	// Make a tempvar so we can replace 0s and 1s later
	tempvar dw
	qui gen `dw' = `var'

	// Do the inverse.
	qui replace `dw' = `dw' ^ (1 / `gamma')

	// Hard replace 0s with 0s and 1s with 1s
	qui replace `dw' = 0 if `var' == 0
	qui replace `dw' = 1 if `var' == 1

	// Overwrite the initial value
	qui replace `var' = `dw'
end

// Inverse-S doesn't have an analytical solution, so we use Newton's root
// finding method to analytically solve for inverse
program define inverse_s_inverse
	args dw gamma

	// I picked this tolerence pretty arbitrarily, but it's really close
	local ee  = 0.0000001
	// Number of iterations. In some testing, convergence is found with only ~26
	// iterations. However, this method is far quicker than a naive grid search,
	// and much more accurate, so we can be a bit conservative and use more
	// iterations than strictly necessary.
	local iter = 50

	// Make a tempvar so we can replace 0s and 1s later
	tempvar  dw0 prob eprob fn0 fn1 grad
	// This is the value of the Inverse-S function
	qui gen `dw0'  = `dw'
	qui gen `fn0'  = 0
	qui gen `fn1'  = 0
	qui gen `grad' = 0

	qui gen `prob' = 0
	qui gen `eprob' = 0.5

	//  Newton's Method for root finding
	forvalues i = 1 / `iter' {
		// Calculate the base
		// eprob is restricted to (0, 1)
		// We want to find the value of prob that makes this function closest to 0
		qui replace `eprob' = 1 / (1 + exp(-`prob'))
		qui replace `fn0' = ((`eprob'^`gamma' / (`eprob'^`gamma' + (1 - `eprob')^`gamma')^(1 / `gamma')) - `dw0')^2

		// Calculate the base + ee
		qui replace `eprob' = 1 / (1 + exp(-(`prob' + `ee')))
		qui replace `fn1'   = ((`eprob'^`gamma' / (`eprob'^`gamma' + (1 - `eprob')^`gamma')^(1 / `gamma')) - `dw0')^2

		// Calculate the gradient of fn
		qui replace `grad' = (`fn1' - `fn0') / `ee'

		// Adjust prob assuming the gradient can tell us where the value of 0 is
		qui replace `prob' = `prob' - (`fn0' / `grad')
	}

	// Recover the decision weight from this
	qui replace `dw0' = 1 / (1 + exp(-(`prob')))

	// Hard replace:
	//   missing values with 0,
	//   0s with 0s ,
	//   and 1s with 1s
	qui replace `dw0' = 0 if `dw0' == .
	qui replace `dw0' = 0 if `dw'  == 0
	qui replace `dw0' = 1 if `dw'  == 1

	// Replace the value passed in with the calculated inverse
	qui replace `dw' = `dw0'
end

// CRRA Utility function prime
program define crra_prime
	args newval outcome r
	qui replace `newval' = `outcome' ^ -`r'
end

// Power Utility function prime
program define power_prime
	args newval outcome r
	qui replace `newval' = `r' * `outcome' ^ (`r' - 1)
end

// CARA Utility function prime
program define cara_prime
	args newval outcome alpha
	qui replace `newval' = exp(-`alpha' * `outcome')
end

// Expo Power Utility function prime
program define expo_prime
	args newval outcome s alpha
	qui replace `newval' = exp(-`alpha' * (`outcome'^(1 - `s'))) * (1 - `s') * `outcome'^(-`s')
end
